package org.zjt.agent.generate.bytebuddy;

import lombok.extern.slf4j.Slf4j;
import net.bytebuddy.ByteBuddy;
import net.bytebuddy.description.type.TypeDescription;
import net.bytebuddy.dynamic.ClassFileLocator;
import net.bytebuddy.dynamic.DynamicType;
import net.bytebuddy.dynamic.loading.ClassLoadingStrategy;
import net.bytebuddy.implementation.FixedValue;
import net.bytebuddy.implementation.MethodDelegation;
import net.bytebuddy.implementation.bind.annotation.*;
import net.bytebuddy.matcher.ElementMatchers;
import net.bytebuddy.pool.TypePool;
import org.zjt.agent.ClassDefineDescription;
import org.zjt.agent.Order;
import org.zjt.agent.bytebuddy.binder.ToStringInterceptor;
import org.zjt.agent.generate.ClassHandler;
import org.zjt.agent.generate.GeneratorEnum;
import org.zjt.agent.generate.MethodRedefineStrategy;
import sun.misc.Unsafe;

import java.lang.reflect.Field;
import java.lang.reflect.Method;
import java.security.PermissionCollection;
import java.security.Permissions;
import java.security.ProtectionDomain;
import java.util.*;
import java.util.concurrent.Callable;

import static net.bytebuddy.matcher.ElementMatchers.named;

/**
 * Description:
 *
 * @author juntao.zhang
 * Date: 2018-07-05 下午7:11
 * @see
 */
@Slf4j
public class ThreadPoolClassHandler implements ClassHandler<ClassDefineDescription> {


    private TypePool typePool = TypePool.Default.ofClassPath();

    private ByteBuddy byteBuddy = new ByteBuddy();

    private Unsafe unsafe;



    public ThreadPoolClassHandler(){
        try {

            Field f = Unsafe.class.getDeclaredField("theUnsafe");
            f.setAccessible(true);
            unsafe = (Unsafe) f.get(null);
        }catch (Exception e){
            throw new RuntimeException(e);
        }
    }


    @Override
    public Optional<byte[]> redefine(ClassDefineDescription classDefine) {
        try {
            DynamicType.Unloaded<Object> unloaded = byteBuddy
                    .rebase(typePool.describe(classDefine.getClassName()).resolve(), ClassFileLocator.ForClassLoader.ofClassPath())
                    .method(named("execute").or(ElementMatchers.named("beforeExecute")).or(ElementMatchers.named("beforeExecute")))
                    .intercept(MethodDelegation.withDefaultConfiguration().withBinders(Morph.Binder.install(OverrideCallable.class)).filter(ElementMatchers.named("callableIntercepte"))
                            .to(ToStringInterceptor.class)).make();
            byte[] bytes = unloaded.getBytes();
            Set<Map.Entry<TypeDescription, byte[]>> entries1 = unloaded.getAuxiliaryTypes().entrySet();
            for (Map.Entry<TypeDescription, byte[]> entry :entries1) {
                log.info("Auxiliary class generate :{}",entry.getKey().getName());
                byte[] value = entry.getValue();
                PermissionCollection perms = new Permissions();
                perms.setReadOnly();
                ProtectionDomain protectionDomain = new ProtectionDomain(null, perms);
                unsafe.defineClass(entry.getKey().getName(),value,0,value.length,ClassLoader.getSystemClassLoader(),protectionDomain);
            }
            log.info("class generate :{} ",classDefine.getClassName());
            return Optional.of(bytes);
        }catch (Exception e){
            log.error("ThreadPoolExecutor class redefine exception",e);
        }

        return Optional.empty();
    }

    @Override
    public boolean support(ClassDefineDescription classDefine) {
        boolean result = Objects.equals("java.util.concurrent.ThreadPoolExecutor", classDefine.getClassName()) ||
                Objects.equals("java.util.concurrent.ScheduledThreadPoolExecutor", classDefine.getClassName())
        || classDefine.getClassName().startsWith("java.util.concurrent.ThreadPoolExecutor$auxiliary");
        return result;
    }

    @Override
    public int order() {
        return Order.HIGH.getValue();
    }

    @Override
    public GeneratorEnum getGenerator() {
        return GeneratorEnum.BYTE_BUDDY;
    }

    @Override
    public void addMethodRedefineStrategies(List<MethodRedefineStrategy> methodRedefineStrategies) {

    }

    @Override
    public void addMethodRedefineStrategie(MethodRedefineStrategy methodRedefineStrategie) {

    }


    public static class MyRunnable implements Runnable{

        private final Runnable runnable;

        public MyRunnable(Runnable runnable) {
            this.runnable = runnable;
        }

        @Override
        public void run() {
            runnable.run();
        }
    }



    @Slf4j
    public static class ThreadPoolExecutorInterceptor{

        /**
         *
         * @param call
         * @param method
         * @param arguments
         * @return
         * @throws Exception
         * @see java.util.concurrent.ThreadPoolExecutor#execute  public void execute(Runnable command)
         * @see java.util.concurrent.ThreadPoolExecutor#beforeExecute  protected void beforeExecute(Thread t, Runnable r)
         * @see java.util.concurrent.ThreadPoolExecutor#afterExecute  protected void afterExecute(Thread t, Runnable r)
         */
        @RuntimeType
        public Object intercept(@This Object object ,@SuperCall Callable<?> call, @Origin Method method, @AllArguments Object[] arguments) throws Exception {
            Object result ;

            log.info("intercept  method：{}",method.getName());
            switch (method.getName()){
                case "execute" :
                    Runnable runnable = Runnable.class.cast(arguments[0]);
                    MyRunnable myRunnable = new MyRunnable(runnable);
                    result = method.invoke(object,myRunnable);
                    break;

                case "beforeExecute":

                    result = call.call();
                    break;

                case "afterExecute":

                    result = call.call();
                    break;
                default:
                    result = call.call();
            }
            return result;
        }
    }
}
